#define BUFFERSIZE 1024
#define SENDSIZE 1040
#include <stdio.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <string.h>
#include <getopt.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <stdlib.h>
#include <fcntl.h>
#include <fstream>
#include <iostream>
#include <botan/botan.h>
#include <malloc.h>
#include "RSA.h"
using namespace Botan;
using namespace std;

byte * encryption(RSA *rsa, byte *text, u32bit *len)
{
  PK_Encrypting_Key* pkkey = dynamic_cast<PK_Encrypting_Key*>(rsa->privateKey);
  if(!pkkey) {
    cout << "no key\n";
  }

  SecureVector<byte> ciphertext;
  rsa->encrypt(text,sizeof(text),pkkey,&ciphertext);

  Pipe pipe;
  pipe.process_msg(ciphertext);

  *len = pipe.remaining(0);

  byte *hand_shake_buf;
  hand_shake_buf = (byte *)malloc(*len);

  cout << "cipher len: " << *len << endl;
  pipe.read(hand_shake_buf, *len);

  return(hand_shake_buf);
}

byte * decryption(RSA* rsa, byte *hand_shake_buf, int recv_size)
{
  SecureVector<byte> ciphertext(hand_shake_buf,recv_size);

  PK_Decrypting_Key* prkey = dynamic_cast<PK_Decrypting_Key*>(rsa->privateKey);
  if(!prkey) {
    cout<<"private key error"<<endl;
  }

  SecureVector<byte> plaintext;
  rsa->decrypt(ciphertext,prkey,&plaintext);

  Pipe pipe;
  pipe.process_msg(plaintext);

  recv_size = pipe.remaining(0);

  free(hand_shake_buf);
  hand_shake_buf = (byte *)malloc(recv_size);
  memset(hand_shake_buf,0,recv_size);
  pipe.read(hand_shake_buf,recv_size);

  return(hand_shake_buf);
}

void get_file_size(int fp, struct stat *file_size)
{
  if(fstat(fp, file_size) == -1) {
    perror("Error: invalid file size\n");
  }

  printf("File size: %d bytes\n\n", (int)file_size->st_size);
}

int main(int argc, char **argv)
{
  char ch, *server_prv_key_file = (char*) "serverprvkey.sra";

  while ((ch = getopt(argc, argv, "f:")) != -1) {
    switch (ch) {
      case 'f':
        server_prv_key_file = optarg;
        break;
      case '?':
        if (optopt == 'f') {
          fprintf(stderr, "Error: option -%c requires an argument\n", optopt);
        } else {
          fprintf(stderr, "Error: unknown option\n");
        }
        exit(0);
    }
  }

  struct sockaddr_in client_addr, server_addr;
  int sock, ssock, reuse_addr=1;

  if ((sock = socket(AF_INET, SOCK_STREAM, 0)) == -1) {
    fprintf(stderr, "Error: cannot open socket\n");
    exit(EXIT_FAILURE);
  }

  setsockopt(sock, SOL_SOCKET, SO_REUSEADDR, &reuse_addr, sizeof(reuse_addr));

  memset((char *) &client_addr, 0, sizeof(client_addr));
  memset((char *) &server_addr, 0, sizeof(server_addr));
  server_addr.sin_family = AF_INET;
  server_addr.sin_port = htons(3344);
  server_addr.sin_addr.s_addr = htonl(INADDR_ANY);

  if(bind(sock, (struct sockaddr*)&server_addr, sizeof(server_addr)) < 0) {
    fprintf(stderr, "Error: cannot bind socket\n");
    exit(EXIT_FAILURE);
  }

  char send_buf[SENDSIZE];
  int fp, recv_size=0;
  size_t client_addr_size = sizeof(struct sockaddr_in);
  LibraryInitializer init;

  listen(sock, 0);
  ssock=accept(sock, (struct sockaddr *)&client_addr, &client_addr_size);
  if (ssock < 0) {
    perror("accept");
  }

  byte *hand_shake_buf;
  hand_shake_buf = (byte *)malloc(64);

  RSA* rsa = new RSA();
  string pass_phrase = "azadeh";
  rsa->loadPrivateKey("rsapriv.pem",pass_phrase);

  if ((recv_size=recv(ssock,hand_shake_buf, 64, 0)) < 0) {
    perror("recv");
  }
  hand_shake_buf = decryption(rsa, hand_shake_buf, recv_size);
  string symkey((char*)hand_shake_buf);
  cout<<"Symmetric key: "<<symkey<<endl;
  free(hand_shake_buf);
  hand_shake_buf = (byte *)malloc(recv_size);
  memset(hand_shake_buf,0,recv_size);

  if ((recv_size=recv(ssock,hand_shake_buf, 64, 0)) < 0) {
    perror("recv");
  }
  hand_shake_buf = decryption(rsa, hand_shake_buf, recv_size);
  string autkey((char*)hand_shake_buf);
  cout<<"Authentication key: "<<autkey<<endl<<endl;
  free(hand_shake_buf);
  hand_shake_buf = (byte *)malloc(recv_size);
  memset(hand_shake_buf,0,recv_size);

  if ((recv_size=recv(ssock,hand_shake_buf, 64, 0)) < 0) {
    perror("recv");
  }
  hand_shake_buf = decryption(rsa, hand_shake_buf, recv_size);
  string send_file((char*)hand_shake_buf);
  cout<<"Ready to send file: "<<send_file<<endl;
  free(hand_shake_buf);

  if ((fp=open(send_file.c_str(), O_RDONLY)) <= 0) {
    fprintf(stderr, "Error: cannot open file %s \n",send_file.c_str());
    exit(EXIT_FAILURE);
  }
  struct stat file_size;
  get_file_size(fp,&file_size);
  close(fp);

  int send_size = 0, cum_send_size = 0;
  fstream filestr;

  filestr.open(send_file.c_str(),ios::in|ios::binary);

  //AES Setup
  string algo="AES";
  string passphrase="you never know";

  if(!have_block_cipher(algo)) {
    std::cout << "Don't know about the block cipher \"" << algo << "\"\n";
    return 1;
  }

  const u32bit key_len = max_keylength_of(algo);
  const u32bit iv_len = block_size_of(algo);

  std::auto_ptr<S2K> s2k(get_s2k("PBKDF2(SHA-1)"));
  s2k->set_iterations(8192);

  SymmetricKey bc_key = s2k->derive_key(key_len, symkey.c_str());//"BLK" + passphrase);
  InitializationVector iv = s2k->derive_key(iv_len, symkey.c_str());//"IVL" + passphrase);
  SymmetricKey mac_key = s2k->derive_key(16, autkey.c_str());//"MAC" + passphrase);
  byte encrypt_buffer[SENDSIZE]={0,};
  Keyed_Filter *cast, *hmac;

  Pipe pipe(cast=get_cipher(algo + "/CBC/PKCS7", bc_key, iv, ENCRYPTION)
              );
  pipe.start_msg();
  pipe.write((byte*)&file_size.st_size,sizeof(file_size.st_size));
  pipe.end_msg();
  //cout<<"size of st_size"<<sizeof(file_size.st_size)<<endl;
  //cout<<"Size of encrypt file_size"<<pipe.remaining()<<endl;
  pipe.read((byte*)send_buf,pipe.remaining());
  send(ssock, send_buf,16, 0);


  while (cum_send_size != file_size.st_size) {
    Pipe pipe(new Fork(
                       hmac=new MAC_Filter("HMAC(SHA-1)", mac_key),
                       cast=get_cipher(algo + "/CBC/PKCS7", bc_key, iv, ENCRYPTION)
                       )
             );
    filestr.read(send_buf, BUFFERSIZE);
    pipe.start_msg();
    pipe.write((byte*)send_buf,BUFFERSIZE);
    pipe.end_msg();
    pipe.read((byte *)encrypt_buffer,SENDSIZE,1);
    send_size = send(ssock, encrypt_buffer, SENDSIZE, 0);
    pipe.read((byte *)encrypt_buffer,20,0);
    send(ssock, encrypt_buffer, 20, 0);
    printf("Read from file: %4d bytes,  Send to client: %4d bytes\n",filestr.gcount(),send_size);
    cum_send_size += filestr.gcount();
  }

  cout<<endl<<"Complete sending file: "<<send_file<<endl;
  filestr.close();
  return 0;
}
